#include <iostream>
#include <string>
#include <vector>

#include "ns3/core-module.h"
#include "ns3/network-module.h"
#include "ns3/applications-module.h"
#include "ns3/bridge-module.h"
#include "ns3/csma-module.h"
#include "ns3/internet-module.h"
#include "ns3/point-to-point-module.h"

NS_LOG_COMPONENT_DEFINE ("tcp-incast");

size_t const nserver_limit = 44;
/* int const nsinker_limit = 44; */
std::string const sw_csma_datarate = "10Mbps";
std::string const server_sw_csma_datarate = "10Mbps";
std::string const sinker_sw_csma_datarate = "10Mbps";
std::string const sw_csma_delay = "5ms";
std::string const server_sw_csma_delay = "5ms";
std::string const sinker_sw_csma_delay = "5ms";

std::string node_name(std::string const& prefix, size_t num);
std::string ipv4_address(size_t x);
void init_csmalink(ns3::CsmaHelper & ch,
        std::string const& datarate, std::string const& delay);
/* void init_vector_csmalink(std::vector<ns3::CsmaHelper> & links, */
/*         std::string const& datarate, std::string const& delay); */

void build_topo_and_run(size_t n);
void worker(size_t n);

int main()
{
    /* ns3::Config::SetDefault ("ns3::OnOffApplication::PacketSize", UintegerValue (210)); */
    /* ns3::Config::SetDefault ("ns3::OnOffApplication::DataRate", StringValue ("448kb/s")); */

    size_t nserver = 3;
    worker(nserver);
}

std::string node_name(std::string const& prefix, size_t num)
{
    return prefix + std::to_string(num);
}

std::string ipv4_address(size_t x)
{
    return "192.168.1." + std::to_string(x);
}

void init_csmalink(ns3::CsmaHelper & ch,
        std::string const& datarate, std::string const& delay)
{
    ch.SetChannelAttribute("DataRate", ns3::StringValue(datarate));
    ch.SetChannelAttribute("Delay", ns3::StringValue(delay));
}

/* void init_vector_csmalink(std::vector<ns3::CsmaHelper> & links, */
/*         std::string const& datarate, std::string const& delay) */
/* { */
/*     for (auto & link : links) init_csmalink(link, datarate, delay); */
/* } */

void build_topo_and_run(size_t n)
{
    ns3::LogComponentEnable("tcp-incast", ns3::LOG_LEVEL_INFO);

    // create nodes
    std::vector<ns3::Ptr<ns3::Node>> servers;
    std::vector<ns3::Ptr<ns3::Node>> sinkers;
    std::vector<ns3::Ptr<ns3::Node>> sw;
    for (size_t i = 0; i < n; i++) servers.emplace_back(ns3::CreateObject<ns3::Node>());
    for (size_t i = 0; i < n; i++) sinkers.emplace_back(ns3::CreateObject<ns3::Node>());
    for (size_t i = 0; i < 2; i++) sw.emplace_back(ns3::CreateObject<ns3::Node>());

    // give the nodes names
    for (size_t i = 0; i < n; i++)
        ns3::Names::Add(node_name("server", i + 1), servers[i]);
    for (size_t i = 0; i < n; i++)
        ns3::Names::Add(node_name("sinker", i + 1), sinkers[i]);
    for (size_t i = 0; i < 2; i++)
        ns3::Names::Add(node_name("switch", i + 1), sw[i]);

    // create csma links
    ns3::CsmaHelper server_sw_csma;
    ns3::CsmaHelper sinker_sw_csma;
    ns3::CsmaHelper sw_csma;

    init_csmalink(server_sw_csma, sw_csma_datarate, sw_csma_delay);
    init_csmalink(sinker_sw_csma, sw_csma_datarate, sw_csma_delay);
    init_csmalink(sw_csma, sw_csma_datarate, sw_csma_delay);

    // netdevice
    std::vector<ns3::NetDeviceContainer> server_sw_link(n);
    std::vector<ns3::NetDeviceContainer> sinker_sw_link(n);
    ns3::NetDeviceContainer sw_link;

    for (size_t i = 0; i < n; i++)
        server_sw_link[i] = server_sw_csma
            .Install(ns3::NodeContainer(servers[i], sw[0]));
    for (size_t i = 0; i < n; i++)
        sinker_sw_link[i] = sinker_sw_csma
            .Install(ns3::NodeContainer(sinkers[i], sw[1]));
    sw_link = sw_csma.Install(ns3::NodeContainer(sw[0], sw[1]));

    // switch netdevice
    std::vector<ns3::NetDeviceContainer> swnd(2);
    for (size_t i = 0; i < 2; i++) {
        for (size_t j = 0; j < n; j++)
            if (i == 0) swnd[i].Add(server_sw_link[j].Get(1));
            else swnd[i].Add(sinker_sw_link[j].Get(1));
        swnd[i].Add(sw_link.Get(i));
    }

    // install bridging code on each switch
    ns3::BridgeHelper bridge;
    for (size_t i = 0; i < 2; i++)
        bridge.Install(sw[i], swnd[i]);

    // install the internet stack (TCP/IP)
    ns3::InternetStackHelper ipstack;
    ns3::NodeContainer endpoints;
    for (size_t i = 0; i < n; i++) {
        endpoints.Add(servers[i]);
        endpoints.Add(sinkers[i]);
    }
    ipstack.Install(endpoints);

    // assign ip addresses
    ns3::NetDeviceContainer ip_devices;
    for (size_t i = 0; i < n; i++) ip_devices.Add(server_sw_link[i].Get(0));
    for (size_t i = 0; i < n; i++) ip_devices.Add(sinker_sw_link[i].Get(0));

    ns3::Ipv4AddressHelper ipv4;
    ipv4.SetBase("192.168.1.0", "255.255.255.0");
    /* ns3::Ipv4InterfaceContainer interfaces = */
    ipv4.Assign(ip_devices);

    // debug ipv4 address
    /* for (auto i = 0u; i < n; i++) { */
    /*     ns3::Ptr<ns3::Ipv4> tmp = servers[i]->GetObject<ns3::Ipv4>(); */
    /*     ns3::Ipv4InterfaceAddress iaddr = tmp->GetAddress(1,0); */
    /*     ns3::Ipv4Address addri = iaddr.GetLocal(); */
    /*     NS_LOG_INFO(addri); */
    /* } */

    /* for (auto i = 0u; i < n; i++) { */
    /*     ns3::Ptr<ns3::Ipv4> tmp = sinkers[i]->GetObject<ns3::Ipv4>(); */
    /*     ns3::Ipv4InterfaceAddress iaddr = tmp->GetAddress(1,0); */
    /*     ns3::Ipv4Address addri = iaddr.GetLocal(); */
    /*     NS_LOG_INFO(addri); */
    /* } */

    /* // create sinkers */
    int port = 50000;
    std::vector<ns3::PacketSinkHelper> sink_helpers;
    for (size_t i = 0; i < n; i++) {
        ns3::Address sink_address(ns3::InetSocketAddress(ns3::Ipv4Address(ipv4_address(i + 1).c_str()), port));
        sink_helpers.emplace_back("ns3::TcpSocketFactory", sink_address);
    }
    std::vector<ns3::ApplicationContainer> sink_apps(n);
    for (size_t i = 0; i < n; i++) {
        sink_apps[i] = sink_helpers[i].Install(sinkers[i]);
        sink_apps[i].Start(ns3::Seconds(1.0));
        sink_apps[i].Stop(ns3::Seconds(60.0));
    }

    // create OnOff apps to send TCP to the sinkers
    std::vector<ns3::OnOffHelper> client_helpers;
    for (size_t i = 0; i < n; i++)
        client_helpers.emplace_back("ns3::TcpSocketFactory", ns3::Address());
    for (auto & i : client_helpers) {
        i.SetAttribute("OnTime", ns3::StringValue("ns3::ConstantRandomVariable[Constant=1]"));
        i.SetAttribute("OffTime", ns3::StringValue("ns3::ConstantRandomVariable[Constant=0]"));
        /* i.SetAttribute("Interval", ns3::TimeValue(ns3::Seconds(0.005))); */
        /* i.SetAttribute("PacketSize", ns3::UintegerValue(1000)); */
    }

    std::vector<ns3::ApplicationContainer> client_apps(n);
    for (size_t i = 0; i < n; i++) {
        ns3::AddressValue remote_address(ns3::InetSocketAddress(ns3::Ipv4Address(ipv4_address(i + n + 1).c_str()), port));
        client_helpers[i].SetAttribute("Remote", remote_address);

        client_apps[i].Add(client_helpers[i].Install(servers[i]));

        client_apps[i].Start(ns3::Seconds(1.0));
        client_apps[i].Stop(ns3::Seconds(60.0));
    }

    /* // test udp */

    /* ns3::UdpEchoServerHelper server (9); */

    /* ns3::ApplicationContainer serverApp = server.Install (servers[0]); */
    /* serverApp.Start (ns3::Seconds (0.5)); */
    /* serverApp.Stop  (ns3::Seconds (60)); */

    /* // ------------------------------------------------------------------ */
    /* // Install multi-switch UDP echo client on t2 */
    /* // ------------------------------------------------------------------ */

    /* ns3::Time     interPacketInterval = ns3::Seconds (0.005); */
    /* uint32_t packetSize          = 1000; */
    /* uint32_t maxPacketCount      = (60 - 2.0) / 0.005; */

    /* ns3::UdpEchoClientHelper client (ns3::Ipv4Address ("192.168.1.4"), 9); */

    /* client.SetAttribute ("MaxPackets", ns3::UintegerValue (maxPacketCount)); */
    /* client.SetAttribute ("Interval",   ns3::TimeValue     (interPacketInterval)); */
    /* client.SetAttribute ("PacketSize", ns3::UintegerValue (packetSize)); */

    /* ns3::ApplicationContainer clientApp = client.Install (sinkers[0]); */
    /* clientApp.Start (ns3::Seconds (0.5)); */
    /* clientApp.Stop  (ns3::Seconds (60)); */

    // configure tracing
    ns3::AsciiTraceHelper ascii;
    server_sw_csma.EnableAsciiAll(ascii.CreateFileStream("trace/tcp.tr"));
    server_sw_csma.EnablePcapAll("trace/tcp");

    sinker_sw_csma.EnableAsciiAll(ascii.CreateFileStream("trace/tcp.tr"));
    sinker_sw_csma.EnablePcapAll("trace/tcp");

    sw_csma.EnableAsciiAll(ascii.CreateFileStream("trace/tcp.tr"));
    sw_csma.EnablePcapAll("trace/tcp");

    /* // another way to configure */
    /* for (size_t i = 0; i < n; i++) */
    /*     server_sw_csma.EnablePcap("trace/server_sw" + std::to_string(i) + ".pcap", ip_devices.Get(i), true, true); */
    /* for (size_t i = 0; i < n; i++) */
    /*     sinker_sw_csma.EnablePcap("trace/sinker_sw" + std::to_string(i) + ".pcap", ip_devices.Get(n + i), true, true); */

    ns3::Simulator::Stop(ns3::Seconds(60));
    ns3::Simulator::Run();
    ns3::Simulator::Destroy();
    NS_LOG_INFO("Done.");
}

void worker(size_t n)
{
    // init topo
    build_topo_and_run(n);
}

